#!/usr/bin/perl -wc

package RFCMail::DNS::AsyncQuery;

use strict;
use base qw(Exporter);
use Carp;
use Errno;
use IO::Socket::INET;
use Event;
use Net::DNS;
use RFCMail::Util::Callback qw(callback);


BEGIN {
    our @EXPORT_OK = qw(async_dns_query
			async_dns_multi_query
			async_dns_query_udp
			async_dns_query_tcp
			async_dns_query_existing_tcp);
}


INIT {
    our $DEFAULT_PORT = 53; # the standard DNS port
    our $DEFAULT_TCP_TIMEOUT = 120;
}


sub async_dns_query {
    my ($callback, @query) = @_;
    async_dns_query_udp(sub {
	my $result = shift;
	if ($result and $result->header->tc) {
	    # Response was truncated; try TCP on same server
	    async_dns_query_tcp($callback, $result->answerfrom, @query);
	}
	else {
	    callback($callback, $result);
	}
    }, @query);
}


sub async_dns_multi_query {
    my ($callback, @queries) = @_;
    my $outstanding = scalar(@queries);
    my @results;
    for (@queries) {
	async_dns_query(sub {
	    my $result = shift;
	    push @results, $result
		if $result;
	    --$outstanding
		or callback($callback, @results);
	}, @$_);
    }
}


sub async_dns_query_udp {
    my $callback = shift;
    my $query = Net::DNS::Packet->new(@_)
	or croak("unable to generate query packet");
    RFCMail::DNS::AsyncQuery::UDP::start($callback, $query);
}


sub async_dns_query_tcp {
    my $callback = shift;
    my $ns = shift;
    our ($DEFAULT_PORT);
    my $sock = IO::Socket::INET->new(PeerAddr => $ns, PeerPort => $DEFAULT_PORT,
				     Proto => "tcp", Blocking => 0);
    async_dns_query_existing_tcp(sub {
	my $result = shift;
	$sock->close;
	# discard result if it is an exception other than "NXDOMAIN"
	if ($result) {
	    my $rcode = $result->header->rcode;
	    undef $result
		if $rcode ne 'NOERROR' and $rcode ne 'NXDOMAIN';
	}
	callback($callback, $result);
    }, $sock, @_);
}


# This procedure is hideously long; I'd like to refine it.

sub async_dns_query_existing_tcp {
    my $callback = shift;
    my $sock = shift;
    our ($DEFAULT_TCP_TIMEOUT);
    my $packet = Net::DNS::Packet->new(@_)
	or croak("unable to generate query packet");
    my $data = pack("n/a*", $packet->data);
    # done with $packet for now, but we'll assemble the response packet in it later
    undef $packet;
    my $event;
    Event->io(fd => $sock, poll => 'w', timeout => $DEFAULT_TCP_TIMEOUT, cb => sub {
	$event = shift;
	if ($event->got) {
	    # ready to write
	    my $wrote = $sock->syswrite($data);
	    if (defined($wrote)) {
		# cut the bit we sent from $data
		substr($data, 0, $wrote, "");
		# return if more to send
		return
		    if $data;
	    }
	    else {
		# some kind of error; just return if it's a blocking issue
		return
		    if $!{EAGAIN} || $!{EWOULDBLOCK};
	    }
	}
	# reaching this point means a timeout, a fatal error, or successful transmission
	# no more write-watcher events in any case
	$event->w->cancel;
	# untransmitted data means we failed
	if ($data) {
	    callback($callback, undef);
	    return;
	}
	# success, on the other hand, means we read back into the now-empty $data
	Event->io(fd => $sock, poll => 'r', timeout => $DEFAULT_TCP_TIMEOUT, cb => sub {
	    $event = shift;
	    if ($event->got) {
		# ready to read; compute the remaining bytes wanted
		# note that we first want two bytes to determine the size, then that many more
		my $have = length($data);
		my $want = ($have < 2 ? 2 : unpack("n", $data) + 2) - $have;
		my $got = $sock->sysread($data, $want, $have);
		if (defined($got)) {
		    # some data has been appended to $data; return if not complete
		    return
			if length($data) < 2;
		    return
			if length($data) < unpack("n", $data) + 2;
		    # we seem to have a winner; extract the packet
		    $data = unpack("n/a", $data);
		    $packet = Net::DNS::Packet->new(\$data);
		    undef $data;
		}
		else {
		    # some kind of error; just return if it's a blocking issue
		    return
			if $!{EAGAIN} || $!{EWOULDBLOCK};
		}
	    }
	    # reaching this point means a timeout, a fatal error, or successful reception
	    # these are all terminal cases, and $packet is correct for all cases
	    $event->w->cancel;
	    callback($callback, $packet);
	});
    });
}



# This sub-package handles transmission of a specific UDP query-in-progress

package RFCMail::DNS::AsyncQuery::UDP;

use strict;
use IO::Socket::INET;
use Net::DNS;
use Event;
use RFCMail::Util::Callback qw(callback);


INIT {
    our @DEFAULT_NAMESERVERS = Net::DNS::Resolver->new->nameservers;
    our $UDP_BUFFER_SIZE = 1500; # technically, the payload shouldn't exceed 512
    our $DEFAULT_PORT = $RFCMail::DNS::AsyncQuery::DEFAULT_PORT;
    our $DEFAULT_MAXTRANS = 4; # maximum number of transmission attempts
    our $DEFAULT_WAIT = 5; # number of seconds to wait between transmission attemtps
    our (%Blackboard, $Socket, $Watcher); # for queries in progress
}


# initiate a query
sub start {
    my ($callback, $query, %opts) = @_;
    our (@DEFAULT_NAMESERVERS, $DEFAULT_PORT, $DEFAULT_MAXTRANS, $DEFAULT_WAIT);
    our (%Blackboard, $Socket, $Watcher);
    # change query ID if it is already on the blackboard
    my $id = $query->header->id;
    while (exists($Blackboard{$id})) { ++$id; $id %= 0x10000; }
    $query->header->id($id);
    # build query-in-progress object; set default values as needs be
    my $self = bless \%opts, __PACKAGE__;
    $self->{callback} = $callback;
    $self->{query} = $query;
    $self->{nameservers} ||= [ @DEFAULT_NAMESERVERS ];
    $self->{port} ||= $DEFAULT_PORT;
    $self->{maxtrans} ||= $DEFAULT_MAXTRANS;
    $self->{wait} ||= $DEFAULT_WAIT;
    if (not scalar(%Blackboard)) {
	# this will be the first item on the blackboard; create the socket and watcher
	$Socket = IO::Socket::INET->new(Proto => 'udp', Blocking => 0)
	    or die "unable to open socket for DNS queries ($!)";
	$Watcher = Event->io(fd => $Socket, cb => \&read_response);
    }
    $Blackboard{$id} = $self;
    $self->{sender} = Event->timer(after => 0, interval => $self->{wait},
				   cb => [ $self, 'transmit_query' ]);
}


# handle an incoming response on the socket (Event->io driven)
sub read_response {
    my $event = shift;
    my ($data, $from, $answer, $self, $question);
    our ($UDP_BUFFER_SIZE, %Blackboard);
    # receive incoming packet, aborting if any error is encountered
    $from = $event->w->fd->recv($data, $UDP_BUFFER_SIZE, MSG_DONTWAIT)
	or return;
    # convert to dotted-quad format
    $from = inet_ntoa((unpack_sockaddr_in($from))[1]);
    # abort if data doesn't make sense as a DNS message
    $answer = Net::DNS::Packet->new(\$data)
	or return;
    $answer->answerfrom($from);
    # abort if the ID is not listed on the blackboard
    $self = ($Blackboard{$answer->header->id})
	or return;
    # it's a real, expected response, but is it what we asked? abort if not
    $question = ($self->{query}->question)[0]->string;
    for ($answer->question) {
	return if $_->string ne $question;
    }
    if ($answer->header->rcode eq 'NOERROR' or $answer->header->rcode eq 'NXDOMAIN') {
	# successful query: finish up
	$self->finish($answer);
    }
    else {
	# an error! don't ask this unhelpful nameserver again: remove it from the list
	@{$self->{nameservers}} = grep($_ ne $from, @{$self->{nameservers}});
	# coerce a retransmit timeout if that was in response to our most recent transmission
	$self->{sender}->again
	    if $from eq $self->{last_ns};
    }
}


# method to attempt a transmission of the query (or force failure by timeout)
sub transmit_query {
    my ($self) = @_;
    our ($Socket);
    if ($self->{maxtrans} and @{$self->{nameservers}}) {
	--$self->{maxtrans};
	# cycle through nameservers on successive queries
	my $ns = shift @{$self->{nameservers}};
	push @{$self->{nameservers}}, $ns;
	# note that this was our most recent nameserver for optimisation purposes
	$self->{last_ns} = $ns;
	my $ns_addr = pack_sockaddr_in($self->{port}, inet_aton($ns));
	# attempt a send; if it works, it works
	$Socket->send($self->{query}->data, MSG_DONTWAIT, $ns_addr);
    }
    else {
	# we're out of bullets (or targets), so give up
	$self->finish(undef);
    }
}


# method to terminate a query, successfully or not
sub finish {
    my ($self, $response) = @_;
    our (%Blackboard, $Socket, $Watcher);
    # dismantle the sender event generator
    (delete $self->{sender})->cancel;
    # wipe ourself off the blackboard
    delete $Blackboard{$self->{query}->header->id};
    if (not scalar(%Blackboard)) {
	# the blackboard is clear; cancel the watcher and close the socket
	$Watcher->cancel;
	undef $Watcher;
	$Socket->close;
	undef $Socket;
    }
    callback($self->{callback}, $response);
}

1;
